#include "stdafx.h"
#include "ProfilerCallback.h"

ProfilerCallback::ProfilerCallback() :
	m_refCount( 0 )
{
}

#pragma region IUnknown

ULONG ProfilerCallback::AddRef() 
{
	return InterlockedIncrement( &m_refCount );
}

ULONG ProfilerCallback::Release() 
{
	long refCount;

	refCount = InterlockedDecrement( &m_refCount );

	if ( refCount == 0 )
		delete this;

	return refCount;
}

HRESULT ProfilerCallback::QueryInterface( REFIID riid, void **ppInterface )
{
	if ( riid == IID_IUnknown )
		*ppInterface = static_cast<IUnknown *>( this ); 
	else if ( riid == IID_ICorProfilerCallback )
		*ppInterface = static_cast<ICorProfilerCallback *>( this );
	else if ( riid == IID_ICorProfilerCallback2 )
		*ppInterface = static_cast<ICorProfilerCallback2 *>( this );
	else if ( riid == IID_ICorProfilerCallback3 )
		*ppInterface = static_cast<ICorProfilerCallback3 *>( this );
	else
	{
		*ppInterface = NULL;
		return E_NOINTERFACE;
	}

	reinterpret_cast<IUnknown *>( *ppInterface )->AddRef();

	return S_OK;
}

#pragma endregion

#pragma region Factory

HRESULT ProfilerCallback::CreateObject( REFIID riid, void **ppInterface )
{
	HRESULT hr = E_NOINTERFACE;

	*ppInterface = NULL;
	if (   (riid == IID_IUnknown)
		|| (riid == IID_ICorProfilerCallback)
		|| (riid == IID_ICorProfilerCallback2) 
		|| (riid == IID_ICorProfilerCallback3))
	{           
		ProfilerCallback *pProfilerCallback;

		pProfilerCallback = new ProfilerCallback();
		if ( pProfilerCallback != NULL )
		{
			hr = S_OK;

			pProfilerCallback->AddRef();

			*ppInterface = static_cast<ICorProfilerCallback *>( pProfilerCallback );
		}
		else
			hr = E_OUTOFMEMORY;
	}    

	return hr;
}

#pragma endregion

#pragma region Enter/Leave Callbacks

__forceinline void ProfilerCallback::Enter( FunctionID functionID )
{
	//std::ofstream myfile;
	//myfile.open ("enter.txt");
	//myfile <<  "Enter\n";
	//myfile.close();
}

__forceinline void ProfilerCallback::Leave( FunctionID functionID )
{
	//std::ofstream myfile;
	//myfile.open ("leave.txt");
	//myfile <<  "Leave\n";
	//myfile.close();
}

void ProfilerCallback::Tailcall( FunctionID functionID )
{
	//std::ofstream myfile;
	//myfile.open ("tailcall.txt");
	//myfile <<  "Tailcall\n";
	//myfile.close();
}

EXTERN_C void __stdcall EnterStub( FunctionID functionID )
{
	ProfilerCallback::Enter( functionID );
}

EXTERN_C void __stdcall LeaveStub( FunctionID functionID )
{
	ProfilerCallback::Leave( functionID );   
}

EXTERN_C void __stdcall TailcallStub( FunctionID functionID )
{
	ProfilerCallback::Tailcall( functionID );
}

void __declspec( naked ) __stdcall EnterNaked3(FunctionIDOrClientID functionIDOrClientID)
{
	__asm
	{
		push eax
		push ecx
		push edx
		push [esp + 16]
		call EnterStub
		pop edx
		pop ecx
		pop eax
		ret 4
	}
}

void __declspec( naked ) __stdcall LeaveNaked3(FunctionIDOrClientID functionIDOrClientID)
{
	__asm
	{
		push eax
		push ecx
		push edx
		push [esp + 16]
		call LeaveStub
		pop edx
		pop ecx
		pop eax
		ret 4
	}
}

void __declspec( naked ) __stdcall TailcallNaked3(FunctionIDOrClientID functionIDOrClientID)
{
	__asm
	{
		push eax
		push ecx
		push edx
		push [esp + 16]
		call TailcallStub
		pop edx
		pop ecx
		pop eax
		ret 4
	}
}

#pragma endregion

#pragma region ICorProfilerCallback

HRESULT ProfilerCallback::Initialize( IUnknown *pICorProfilerInfoUnk )
{
	HRESULT hr;

	hr = pICorProfilerInfoUnk->QueryInterface( IID_ICorProfilerInfo, (void **)&m_pProfilerInfo );   	
	if (FAILED(hr))
		return hr;

	hr = pICorProfilerInfoUnk->QueryInterface( IID_ICorProfilerInfo2, (void **)&m_pProfilerInfo2 );
	if (FAILED(hr))
		return hr;

	hr = pICorProfilerInfoUnk->QueryInterface( IID_ICorProfilerInfo3, (void **)&m_pProfilerInfo3 );
	if (FAILED(hr))
		return hr;

	// TODO: prepare event mask
	m_dwEventMask = 
		(DWORD) COR_PRF_MONITOR_ENTERLEAVE
		;

	hr = m_pProfilerInfo->SetEventMask( m_dwEventMask );
	if (FAILED(hr))    
		return hr;

	hr = m_pProfilerInfo3->SetEnterLeaveFunctionHooks3( 
		EnterNaked3,
		LeaveNaked3,
		TailcallNaked3 );

	return S_OK;
}

HRESULT ProfilerCallback::Shutdown()
{
	return S_OK;
}

#pragma endregion
